{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import math\n",
    "import time\n",
    "\n",
    "\n",
    "class HMM:\n",
    "    '''\n",
    "    Implement an HMM\n",
    "    \n",
    "    Use notation and formulae from Jurafsky--Martin \n",
    "    (https://web.stanford.edu/~jurafsky/slp3/9.pdf)\n",
    "    \n",
    "    Assume the hidden states are from 1, 2, ..., to N.  State 0 and \n",
    "    State N+1 are reserved from the START and STOP state, which occur \n",
    "    at exactly time 0 and time T+1 of a T long observation sequence.\n",
    "    \n",
    "    Assume observation values are from 1, 2, ..., to some max values \n",
    "    M. Observation 0 and Observation M+1 are reserved and are seen\n",
    "    at exactly time 0 and time T+1 of a T long observation sequence.\n",
    "    '''\n",
    "    \n",
    "    def __init__(self, a, b):\n",
    "        '''\n",
    "        Initialize the HMM with given parameters\n",
    "        \n",
    "        Initialize any missing parameter with random probability\n",
    "        values.\n",
    "        \n",
    "        \"a\", the transmission probabilities, is a (N+2) x (N+2) matrix \n",
    "        such that a(i, j) is the \n",
    "        probability of moving from hidden state i to hidden state j. \n",
    "        Since one cannot move to State 0 from any state, Column 0 \n",
    "        has 0 values.  Similarly Row N+1 has 0 values.\n",
    "        \n",
    "        \"b\", the emission probabilities, is a (N+2) x (M+2) matrix such\n",
    "        that b(i, j) is the probability of emitting output j in hidden\n",
    "        state i.  Since\n",
    "        Output 0 is only emitted in State 0, Row 0 has a 1.0 in \n",
    "        Column 0, and 0 values otherwise.  Similarly, Row N+1 has\n",
    "        a 1.0 in Column M+1 and 0 values otherwise.\n",
    "        '''\n",
    "        \n",
    "        \n",
    "        # Check input\n",
    "        if (a.shape[0] != a.shape[1] or\n",
    "            a.shape[0] != b.shape[0]):\n",
    "            raise Exception(\"Improper a or b matrix shapes.\")\n",
    "            \n",
    "        # N is number of regular states     \n",
    "        self.N = a.shape[0] - 2\n",
    "        # M is number of regular outputs\n",
    "        self.M = b.shape[1] - 2\n",
    "        \n",
    "        self.a = a\n",
    "        self.b = b\n",
    "        \n",
    "        # Useful macros\n",
    "        self.START = 0 # start state\n",
    "        self.FINISH = self.N+1 # finish state\n",
    "        self.STATES = range(1, self.N+1) # set of regular states\n",
    "        \n",
    "        # placeholders\n",
    "        self.O = [] # observation sequence\n",
    "        self.T = 0 # number of observations\n",
    "        self.alpha_table = None\n",
    "        self.beta_table = None\n",
    "        self.viterbi_table = None\n",
    "        self.viterbi_bpointers = None\n",
    "        \n",
    "    def print(self):\n",
    "        print(f\"N {self.N} M {self.M} T {self.T}\")\n",
    "        print(\"Transmission probabilities 'a'\")\n",
    "        print(self.a)\n",
    "        print(\"Emission probabilities 'b'\")\n",
    "        print(self.b)\n",
    "        \n",
    "    def set_a_b(self, a, b):\n",
    "        \n",
    "        if (a.shape[0] != a.shape[1] or a.shape[0] != b.shape[0] or\n",
    "            a.shape[0] != self.N+2 or b.shape[1] != self.M+2):\n",
    "            raise Exception('Incompatible a, b shapes')\n",
    "        \n",
    "        self.a = a\n",
    "        self.b = b\n",
    "        self.reset_alpha_beta_tables()\n",
    "        \n",
    "    def reset_alpha_beta_tables(self):\n",
    "        \n",
    "        # intialize alpha values to save recursive calls\n",
    "        self.alpha_table = np.array(\n",
    "            [[-1.0] * (self.N+2)] * (self.T+2))\n",
    "        self.alpha_table[self.T+1, :self.N+1] = 0.0\n",
    "        self.alpha_table[0, 0] = 1.0\n",
    "        self.alpha_table[0, 1:] = 0.0\n",
    "        \n",
    "        # initialize beta values\n",
    "        self.beta_table = np.array(\n",
    "            [[-1.0] * (self.N+2)] * (self.T+1)) \n",
    "        # Recall beta(T+1, i) is not defined.\n",
    "        \n",
    "    def set_observations(self, O):\n",
    "        '''\n",
    "        Set the observation sequence\n",
    "        \n",
    "        Assume the first observation is 0 and the last M+1, and\n",
    "        the rest are in [1,...,M]. \n",
    "        '''\n",
    "        \n",
    "        # Check that observations are valid\n",
    "        if (not set(O[1:-1]).issubset(set(range(1, self.M+1))) or\n",
    "            O[0] != 0 or O[-1] != self.M+1):\n",
    "            print(O)\n",
    "            raise Exception(\"Invalid observation sequence.\")\n",
    "        self.O = O\n",
    "        self.T = len(O)-2\n",
    "        self.reset_alpha_beta_tables()\n",
    "        \n",
    "    \n",
    "    def forward(self):\n",
    "        '''Fill alpha_table, i.e., alpha_t(j)'s'''\n",
    "        \n",
    "        # This should be called after set_observation, \n",
    "        # else self.T == 0.\n",
    "        \n",
    "        if (self.T == 0):\n",
    "            raise Exception(\"Premature call to alpha().\")\n",
    "        \n",
    "        ##################################################\n",
    "        # WRITE YOUR CODE HERE ###########################\n",
    "        ##################################################\n",
    "        # Please follow the algorithm as described in \n",
    "        # https://web.stanford.edu/~jurafsky/slp3/9.pdf\n",
    "        # from 9.15 to 9.17\n",
    "        # First fill for t = 0\n",
    "        # Fill remaining values \"recursively\"\n",
    "        \n",
    "        self.alpha_table[0, self.START] = 1.0\n",
    "#         for j in range(1, self.N+1):\n",
    "#             self.alpha_table[0, j] = 0.0\n",
    "        self.alpha_table[0, 1:self.N+1] = 0.0\n",
    "            \n",
    "        # Fill remaining values \"recursively\"\n",
    "        for t in range(1, self.T+2):\n",
    "#             for j in range(self.N+2):\n",
    "#                 self.alpha_table[t, j] = sum(\n",
    "#                     [self.alpha_table[t-1, i] * self.a[i, j]  \n",
    "#                      for i in range(self.N+2)]\n",
    "#                     ) * self.b[j, self.O[t]]\n",
    "            self.alpha_table[t, :self.N+2] = np.matmul(\n",
    "                self.alpha_table[t-1, :], self.a[:, :self.N+2]) * self.b[:self.N+2, self.O[t]]\n",
    "                \n",
    "        \n",
    "    def viterbi(self, t, j):\n",
    "                \n",
    "        if (self.T==0 or t < 1 or t > self.T+1 or j < 0 or\n",
    "            j > self.N + 1 or \n",
    "            (t == self.T+1 and j != self.FINISH)):\n",
    "            raise Exception(\"Invalid call to viterbi().\")\n",
    "            \n",
    "        # Check if value already available\n",
    "        rv = self.viterbi_table[t, j]\n",
    "        if rv <= -1:\n",
    "            if t == 1:\n",
    "                rv = self.a[self.START, j] * self.b[j, self.O[1]]\n",
    "                self.viterbi_bpointers[t, j] = self.START\n",
    "            else:\n",
    "                if t == self.T+1:\n",
    "                    possibles = [(self.viterbi(t-1, i) * \n",
    "                                  self.a[i,j] , i) \n",
    "                                 for i in self.STATES]\n",
    "                else:             \n",
    "                    possibles = [(self.viterbi(t-1, i) * \n",
    "                                  self.a[i,j] * self.b[j, self.O[t]],\n",
    "                                  i)\n",
    "                                 for i in self.STATES]\n",
    "                    \n",
    "                rv, max_i = max(possibles, key=lambda tup: tup[0])\n",
    "                self.viterbi_bpointers[t, j] = max_i\n",
    "            self.viterbi_table[t, j] = rv\n",
    "        return rv\n",
    "    \n",
    "    def decode(self, O=None):\n",
    "        '''Return the most likely hidden state sequenence\n",
    "           for observation sequence obs based \n",
    "           on self.a and self.b, and its probability.\n",
    "           O is the observations'''\n",
    "        \n",
    "        ##################################################\n",
    "        # WRITE YOUR CODE HERE ###########################\n",
    "        ##################################################\n",
    "        # Your code should implement the VITERBI function\n",
    "        # in Figure 9.11 on Jurafsky-Martin \n",
    "        # https://web.stanford.edu/~jurafsky/slp3/9.pdf.\n",
    "        # In particular you should store probabilities in\n",
    "        # a table and store backpointers in another table.\n",
    "        # You need to construct the best hidden state sequence\n",
    "        # using the backpointers and return it along with its\n",
    "        # probability.  For this you will need to use self.a\n",
    "        # and self.b, and values self.T, but you don't need to\n",
    "        # use the alpha or beta tables.\n",
    "        \n",
    "        if O is not None:\n",
    "            self.set_observations(O)\n",
    "        self.viterbi_table = np.ones((self.T+2, self.N+2)) * -1.0\n",
    "        self.viterbi_bpointers = np.ones((self.T+2, self.N+2),\n",
    "                                         dtype = int) * -1\n",
    "        \n",
    "        p = self.viterbi(self.T+1, self.FINISH)\n",
    "        bp_sequence = [self.FINISH]\n",
    "        j = self.FINISH\n",
    "        for t in range(self.T+1,0,-1):\n",
    "                j = self.viterbi_bpointers[t, j]\n",
    "                bp_sequence.append(j)\n",
    "        return (list(reversed(bp_sequence)), p)\n",
    "    \n",
    "    def backward(self):\n",
    "        '''Fill beta_table, i.e., beta_t(i)'s'''\n",
    "               \n",
    "        # Unlike alpha we don't define beta(T+1, i).\n",
    "        # This should be called after set_observation, \n",
    "        # else self.T == 0.\n",
    "        \n",
    "        if (self.T == 0):\n",
    "            raise Exception(\"Premature call to beta()\")\n",
    "            \n",
    "        \n",
    "        ##################################################\n",
    "        # WRITE YOUR CODE HERE ###########################\n",
    "        ##################################################\n",
    "        # Please follow the algorithm as described in \n",
    "        # https://web.stanford.edu/~jurafsky/slp3/9.pdf\n",
    "        # from 9.27 to 9.30\n",
    "        # First fill for t = T\n",
    "        # Fill for t=T-1 to 1 \"recursively\"\n",
    "        # Fill in for t = 0 and i = START\n",
    "        \n",
    "        # First fill for t = T\n",
    "        for i in range(1, self.N+1):\n",
    "            self.beta_table[self.T, i] = self.a[i, self.FINISH]\n",
    "            \n",
    "        # Fill for t=1 to T-1 \"recursively\"\n",
    "        for t in range(self.T-1, 0, -1):\n",
    "#             for i in range(1, self.N+1):\n",
    "#                 self.beta_table[t, i] = sum(\n",
    "#                     [self.a[i, j] * self.b[j,  self.O[t+1]] * \n",
    "#                      self.beta_table[t+1, j] for j in self.STATES])\n",
    "            for i in range(1, self.N+1):\n",
    "                self.beta_table[t, i] = np.sum(\n",
    "                    self.a[i, 1:self.N+1] * self.b[1:self.N+1,  self.O[t+1]] * \n",
    "                     self.beta_table[t+1, 1:self.N+1])\n",
    "        \n",
    "        # Fill in for t = 0 and i = START\n",
    "#         self.beta_table[0, self.START] = sum(\n",
    "#                 [self.a[self.START, j] * self.b[j,  self.O[1]] * \n",
    "#                  self.beta_table[1, j] for j in self.STATES])\n",
    "        self.beta_table[0, self.START] = np.sum(\n",
    "                self.a[self.START, 1:self.N+1] * self.b[1:self.N+1,  self.O[1]] * \n",
    "                 self.beta_table[1, 1:self.N+1])\n",
    "        \n",
    "        \n",
    "        \n",
    "    def xi(self, t, i, j):\n",
    "        '''\n",
    "        Return xi_t(i, j)\n",
    "        '''        \n",
    "\n",
    "        # The basic formula works for t=0 as well. For t=T we need\n",
    "        # a special case because beta(t+1, j) is not defined for all j.\n",
    "                \n",
    "        if self.alpha_table[self.T+1, self.FINISH] <= 0.0:\n",
    "            raise Exception(\"Impossible observation sequence\")\n",
    "    \n",
    "        ##################################################\n",
    "        # WRITE YOUR CODE HERE ###########################\n",
    "        ##################################################\n",
    "        # Please follow formula 9.37 in \n",
    "        # https://web.stanford.edu/~jurafsky/slp3/9.pdf\n",
    "        # to compute the value of xi. \n",
    "        # You may need to take care of some corner cases carefully. \n",
    "        \n",
    "        if t >= 0 and t < self.T:\n",
    "            rv = (self.alpha_table[t, i] * self.a[i, j] * \n",
    "                   self.b[j, self.O[t+1]] * self.beta_table[t+1, j]\n",
    "                 )/ self.alpha_table[self.T+1, self.FINISH]            \n",
    "        elif t == self.T: # non-zero only for j == FINISH\n",
    "            if j != self.FINISH:\n",
    "                rv = 0.0\n",
    "            else:\n",
    "                rv = (self.alpha_table[t, i] * self.a[i, j] / \n",
    "                  self.alpha_table[self.T+1, self.FINISH])\n",
    "        else:\n",
    "            raise Exception(\"Invalid call to xi().\")\n",
    "        \n",
    "        return rv\n",
    "    \n",
    "    def gamma(self, t, j):\n",
    "        '''\n",
    "        Return gamma_t(j)\n",
    "        '''  \n",
    "        \n",
    "        rv = 0\n",
    "        \n",
    "        ##################################################\n",
    "        # WRITE YOUR CODE HERE ###########################\n",
    "        ##################################################\n",
    "        # Your code should compute gamma_t(j) as specified\n",
    "        # in Figure 9.16 of Jurafsky-Martin \n",
    "        # https://web.stanford.edu/~jurafsky/slp3/9.pdf. \n",
    "        # For this you wil need to use self.alpha_table, \n",
    "        # self.beta_table, self.T, and values such as self.START\n",
    "        # (the number corresponding to start state, set at 0),\n",
    "        # and self.FINISH (the number corresponding to the\n",
    "        # the end state, set at self.N+1).\n",
    "        \n",
    "        if self.alpha_table[self.T+1, self.FINISH] <= 0.0:\n",
    "            raise Exception(\"Impossible observation sequence\")\n",
    "        \n",
    "        if t >0 and t < self.T+1:\n",
    "            rv = (self.alpha_table[t, j] * self.beta_table[t, j] / \n",
    "                  self.alpha_table[self.T+1, self.FINISH])\n",
    "        elif t == 0:\n",
    "            if j == self.START:\n",
    "                rv = 1.0\n",
    "            else:\n",
    "                rv = 0.0\n",
    "        elif t == self.T+1:\n",
    "            if j == self.FINISH:\n",
    "                rv = 1.0\n",
    "            else:\n",
    "                rv = 0.0\n",
    "        else:\n",
    "            raise Exception(\"Invalid call to gamma()\")\n",
    "            \n",
    "        return rv\n",
    "\n",
    "    def a_from_xi(self, i, j, rtype='ratio'):\n",
    "        '''\n",
    "        Return new a_{ij} computed from xi\n",
    "        rtype: we define two return types, \"ratio\" or \"separate\", \n",
    "        the separate return type will be used in later functions. \n",
    "        '''\n",
    "    \n",
    "        ##################################################\n",
    "        # WRITE YOUR CODE HERE ###########################\n",
    "        ##################################################\n",
    "        # Your code should compute a(hat) as described\n",
    "        # in Figure 9.16 of Jurafsky-Martin \n",
    "        # https://web.stanford.edu/~jurafsky/slp3/9.pdf. \n",
    "        \n",
    "        if j == self.START or i == self.FINISH:\n",
    "            numerator = 0.0\n",
    "            denominator = 1.0\n",
    "        else:\n",
    "            numerator = sum([self.xi(t, i, j) for t in \n",
    "                                range(0, self.T+1)])\n",
    "                \n",
    "            # Denominator in Jurafsky-Martin can be improved.    \n",
    "            #denominator = sum([self.xi(t, i, k) for t in \n",
    "            #                   range(0, self.T+1) for k in \n",
    "            #               range(self.N+2)])\n",
    "            denominator = sum([self.gamma(t, i) for t in\n",
    "                               range(self.T+1)])\n",
    "\n",
    "        if rtype == 'separate':\n",
    "            rv = (numerator, denominator)\n",
    "        elif math.isclose(denominator, 0.0):\n",
    "            print(\"Divide by zero in i={} j={}\".format(i, j))\n",
    "            rv = np.nan\n",
    "        else:\n",
    "            rv = numerator/denominator\n",
    "        return rv\n",
    "        \n",
    "        return rv\n",
    "        \n",
    "    def b_from_gamma(self, j, v, rtype='ratio'):\n",
    "        '''\n",
    "        Return new b_{ij} computed from xi\n",
    "        rtype: we define two return types, \"ratio\" or \"separate\", \n",
    "        the separate return type will be used in later functions. \n",
    "        '''\n",
    "    \n",
    "        ##################################################\n",
    "        # WRITE YOUR CODE HERE ###########################\n",
    "        ##################################################\n",
    "        # Your code should compute b(hat) as described\n",
    "        # in Figure 9.16 of Jurafsky-Martin \n",
    "        # https://web.stanford.edu/~jurafsky/slp3/9.pdf. \n",
    "        \n",
    "        numerator = sum([self.gamma(t, j) for t in \n",
    "                         range(0, self.T+2) if self.O[t] == v])\n",
    "        denominator = sum([self.gamma(t, j) for t in \n",
    "                           range(0, self.T+2)])\n",
    "\n",
    "        if rtype == 'separate':\n",
    "            rv = (numerator, denominator)    \n",
    "        elif math.isclose(denominator, 0.0):\n",
    "            print(\"Divide by zero in j={}, v={}\".format(j, v))\n",
    "            rv = np.nan\n",
    "        else:\n",
    "            rv = numerator/denominator\n",
    "        return rv\n",
    "    \n",
    "    def new_a(self):\n",
    "        rv = np.array([[-1.0] * (self.N+2)] * (self.N+2))\n",
    "        for i in range(self.N+2):\n",
    "            for j in range(self.N+2):\n",
    "                rv[i,j] = self.a_from_xi(i, j)\n",
    "        return rv\n",
    "    \n",
    "    def new_b(self):\n",
    "        rv = np.array([[-1.0] * (self.M+2)] * (self.N+2))\n",
    "        for j in range(self.N+2):\n",
    "            for v in range(self.M+2):\n",
    "                rv[j,v] = self.b_from_gamma(j,v)\n",
    "        return rv\n",
    "    \n",
    "    def fit(self, O, max_iter=10, verbose=False):\n",
    "        '''\n",
    "        Run the forward-backward algorithm for a single observation\n",
    "        '''\n",
    "        \n",
    "        self.set_observations(O)\n",
    "        converged = False\n",
    "        i = 0\n",
    "        while not converged and i < max_iter:\n",
    "            i += 1\n",
    "            \n",
    "            self.forward()\n",
    "            self.backward()\n",
    "            \n",
    "            new_a = self.new_a()\n",
    "            new_b = self.new_b()\n",
    "            \n",
    "            if verbose:\n",
    "                print(f'Iteration {i}')\n",
    "                print(new_b.T)\n",
    "                print(new_a.T)\n",
    "        \n",
    "            if (np.allclose(new_a, self.a) and\n",
    "                    np.allclose(new_b, self.b)):\n",
    "                converged = True\n",
    "            else:\n",
    "                self.set_a_b(new_a, new_b)\n",
    "        return(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# A toy example for testing your code, based on Jason Eisner's\n",
    "# Workbook\n",
    "\n",
    "transmission_prob = np.array([\n",
    "    [0. , 0.5, 0.5, 0. ],\n",
    "    [0. , 0.8, 0.1, 0.1],\n",
    "    [0. , 0.2, 0.7, 0.1],\n",
    "    [0. , 0. , 0. , 0. ]])\n",
    "\n",
    "emission_prob = np.array([\n",
    "    [1. , 0. , 0. , 0. , 0. ],\n",
    "    [0. , 0.7, 0.2, 0.1, 0. ],\n",
    "    [0. , 0. , 0.3, 0.7, 0. ],\n",
    "    [0. , 0. , 0. , 0. , 1. ]])\n",
    "\n",
    "observation_short = [\n",
    "    0,\n",
    "    2, 3, 3, 2,\n",
    "    4\n",
    "]\n",
    "\n",
    "T_observation_short = len(observation_short) -2\n",
    "\n",
    "observation_long = [\n",
    "    0,\n",
    "    2, 3, 3, 2,\n",
    "    3, 2, 3, 2,\n",
    "    2, 3, 1, 3,\n",
    "    3, 1, 1, 1,\n",
    "    2, 1, 1, 1,\n",
    "    3, 1, 2, 1,\n",
    "    1, 1, 2, 3,\n",
    "    3, 2, 3, 2, \n",
    "    2,\n",
    "    4\n",
    "]\n",
    "\n",
    "T_observation_long = len(observation_long)-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "hmms = HMM(transmission_prob, emission_prob)\n",
    "hmms.set_observations(observation_short)\n",
    "hmms.forward()\n",
    "expected_alpha = np.array(\n",
    "      [[1.        , 0.        , 0.        , 0.        ],\n",
    "       [0.        , 0.1       , 0.15      , 0.        ],\n",
    "       [0.        , 0.011     , 0.0805    , 0.        ],\n",
    "       [0.        , 0.00249   , 0.040215  , 0.        ],\n",
    "       [0.        , 0.002007  , 0.00851985, 0.        ],\n",
    "       [0.        , 0.        , 0.        , 0.00105268]])\n",
    "\n",
    "assert np.allclose(hmms.alpha_table, expected_alpha)\n",
    "\n",
    "hmms.backward()\n",
    "expected_beta = np.array([\n",
    "    [ 0.00105268, -1.        , -1.        , -1.        ],\n",
    "    [-1.        ,  0.0011457 ,  0.0062541 , -1.        ],\n",
    "    [-1.        ,  0.00327   ,  0.01263   , -1.        ],\n",
    "    [-1.        ,  0.019     ,  0.025     , -1.        ],\n",
    "    [-1.        ,  0.1       ,  0.1       , -1.        ]])\n",
    "assert np.allclose(hmms.beta_table, expected_beta)\n",
    "\n",
    "expected_decode = ([0, 2, 2, 2, 2, 3], 0.0007563149999999998)\n",
    "assert hmms.decode(hmms.O) == expected_decode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "hmml = HMM(transmission_prob, emission_prob)\n",
    "hmml.fit(observation_long, max_iter = 20)\n",
    "expected_a = np.array([\n",
    "    [0.00000000e+00, 5.68740363e-15, 1.00000000e+00, 0.00000000e+00],\n",
    "    [0.00000000e+00, 9.33778559e-01, 6.62214406e-02, 1.21039999e-15],\n",
    "    [0.00000000e+00, 7.18667227e-02, 8.64943927e-01, 6.31893502e-02],\n",
    "    [0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]])\n",
    "expected_b = np.array([\n",
    "    [1.        , 0.        , 0.        , 0.        , 0.        ],\n",
    "    [0.        , 0.64048263, 0.14806965, 0.21144771, 0.        ],\n",
    "    [0.        , 0.        , 0.53439047, 0.46560953, 0.        ],\n",
    "    [0.        , 0.        , 0.        , 0.        , 1.        ]])\n",
    "assert np.allclose(hmml.a, expected_a)\n",
    "assert np.allclose(hmml.b, expected_b)\n",
    "\n",
    "expected_decode_result = ([0, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, \n",
    "                           1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, \n",
    "                           2, 2, 2, 2, 2, 2, 3],\n",
    "                          1.4779964597903278e-16)\n",
    "assert np.allclose(hmml.decode(hmml.O)[0], expected_decode_result[0])\n",
    "assert np.isclose(hmml.decode(hmml.O)[1], expected_decode_result[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class HMMMulti(HMM):\n",
    "    '''\n",
    "    Run the forward-backward algorithm for multiple independent observations\n",
    "    \n",
    "    Follow the description in Sec 4.1 of \"Training Hidden Markov Models with \n",
    "    Multiple Observations – A Combinatorial Method\" by Li, Parizeau, \n",
    "    and Plamondon, available at \n",
    "    http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.335.1457&rep=rep1&type=pdf\n",
    "    '''    \n",
    "    \n",
    "    def update_new_a_num_deno(self):\n",
    "        for i in range(self.N+2):\n",
    "            for j in range(self.N+2):\n",
    "                (num, deno) = self.a_from_xi(i, j, 'separate')\n",
    "                self.new_a_num[i,j] += num\n",
    "                self.new_a_deno[i,j] += deno\n",
    "    \n",
    "    def update_new_b_num_deno(self):\n",
    "        for j in range(self.N+2):\n",
    "            for v in range(self.M+2):\n",
    "                (num, deno) = self.b_from_gamma(j, v, 'separate')\n",
    "                self.new_b_num[j, v] += num\n",
    "                self.new_b_deno[j, v] += deno\n",
    "                \n",
    "    def get_likelihoods(self):\n",
    "        '''Get the likelihood at each iteration of\n",
    "           the fit/learning step.\n",
    "           \n",
    "           Assume this is called only after a call to fit().'''\n",
    "        \n",
    "        ##################################################\n",
    "        # WRITE YOUR CODE HERE ###########################\n",
    "        ##################################################\n",
    "        # You have to write code here, and make changes in \n",
    "        # other parts of this \n",
    "        # class as well to store the likelihoods of the \n",
    "        # entire set of observations Oset used in fit,\n",
    "        # at each step of the fit process, and return them\n",
    "        # when this function is called.\n",
    "        \n",
    "        prob = 0 # in log space\n",
    "        self.forward()\n",
    "        prob += np.log(self.alpha_table[\n",
    "            self.T+1, self.FINISH])\n",
    "        return np.exp(prob)\n",
    "\n",
    "\n",
    "\n",
    "    def fit(self, Oset, max_iter=10, verbose=False):\n",
    "        \n",
    "        self.new_a_num = np.array([[0.0] * (self.N+2)] * (self.N+2))\n",
    "        self.new_a_deno = np.array([[0.0] * (self.N+2)] * (self.N+2))\n",
    "        self.new_b_num = np.array([[0.0] * (self.M+2)] * (self.N+2))\n",
    "        self.new_b_deno = np.array([[0.0] * (self.M+2)] * (self.N+2))\n",
    "\n",
    "        \n",
    "        converged = False\n",
    "        i = 0\n",
    "        \n",
    "        while not converged and i < max_iter:\n",
    "            i += 1\n",
    "                   \n",
    "            for O in Oset:\n",
    "                self.set_observations(O)\n",
    "                \n",
    "                self.forward()\n",
    "                self.backward()\n",
    "                \n",
    "                \n",
    "            \n",
    "                self.update_new_a_num_deno()\n",
    "                self.update_new_b_num_deno()\n",
    "            \n",
    "            new_a = self.new_a_num / self.new_a_deno\n",
    "            new_b = self.new_b_num / self.new_b_deno\n",
    "            \n",
    "            if verbose:\n",
    "                print(f'Iteration {i}')\n",
    "                print(new_b.T)\n",
    "                print(new_a.T)\n",
    "            else:\n",
    "                print(f'Iteration {i}')\n",
    "        \n",
    "            if (np.allclose(new_a, self.a) and\n",
    "                    np.allclose(new_b, self.b)):\n",
    "                converged = True\n",
    "            else:\n",
    "                self.set_a_b(new_a, new_b)\n",
    "                self.new_a_num.fill(0)\n",
    "                self.new_a_deno.fill(0)\n",
    "                self.new_b_num.fill(0)\n",
    "                self.new_b_deno.fill(0)\n",
    "                \n",
    "                \n",
    "        return(i)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "hmmml = HMMMulti(transmission_prob, emission_prob)\n",
    "hmmml.fit([observation_long]*5, max_iter = 20)\n",
    "expected_a = np.array([\n",
    "    [0.00000000e+00, 5.68740363e-15, 1.00000000e+00, 0.00000000e+00],\n",
    "    [0.00000000e+00, 9.33778559e-01, 6.62214406e-02, 1.21039999e-15],\n",
    "    [0.00000000e+00, 7.18667227e-02, 8.64943927e-01, 6.31893502e-02],\n",
    "    [0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]])\n",
    "expected_b = np.array([\n",
    "    [1.        , 0.        , 0.        , 0.        , 0.        ],\n",
    "    [0.        , 0.64048263, 0.14806965, 0.21144771, 0.        ],\n",
    "    [0.        , 0.        , 0.53439047, 0.46560953, 0.        ],\n",
    "    [0.        , 0.        , 0.        , 0.        , 1.        ]])\n",
    "assert np.allclose(hmml.a, expected_a)\n",
    "assert np.allclose(hmml.b, expected_b)\n",
    "expected_likelihoods = [8.463524947168683e-21, 6.45760751286262e-19,\n",
    "                        2.604423720662633e-18, 3.772405535111533e-18,\n",
    "                        3.934445572171057e-18, 3.945375545182965e-18,\n",
    "                        3.9459945227193204e-18, 3.946030509476693e-18,\n",
    "                        3.946032758826868e-18, 3.946032909436204e-18,\n",
    "                        3.94603292005471e-18]\n",
    "assert np.allclose(hmmml.get_likelihoods(), expected_likelihoods)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import nltk\n",
    "from nltk.corpus import brown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class POS:\n",
    "    '''\n",
    "        Partition given sentences, or those obtained from the brown\n",
    "        corpus, equally into fit sentences and test sentences.\n",
    "        For each obtain tags using nltk tagger.  Initialize \"a\"\n",
    "        with the tags from all sentences, but with uniform\n",
    "        transition probabilities.  Initialize \"b\" with all words\n",
    "        in all sentences.  But the emission probability \n",
    "        initialization should only use the words in fit sentences.\n",
    "        For each word in the fit sentence, give its most common\n",
    "        tag \"high\" probability of emitting it.  All other words\n",
    "        get a \"low\" probability of being emitted by a tag.  (The\n",
    "        high and low values are different for each tag because of\n",
    "        normalization.)\n",
    "        '''\n",
    "        \n",
    "    def __init__(self, sents=None, category=None, \n",
    "                 num_sents=40,\n",
    "                 b_bound = 1*10**(-4), \n",
    "                 tagset=\"universal\"):\n",
    "        \n",
    "        if sents is None:\n",
    "            self.sents = brown.sents(categories=category)[:num_sents]\n",
    "            self.tagged_sents = brown.tagged_sents(\n",
    "                categories=category, tagset=tagset)[:num_sents]\n",
    "        else:\n",
    "            self.sents = [nltk.word_tokenize(s) for s in sents]\n",
    "            self.tagged_sents = [nltk.pos_tag(s, tagset = tagset)\n",
    "                                 for s in self.sents]\n",
    "            \n",
    "        self.fsize = int(len(self.sents) * 0.5)\n",
    "        self.fit_tsents = self.tagged_sents[:self.fsize]\n",
    "        self.test_tsents = self.tagged_sents[self.fsize:]\n",
    "        tagged_words = [tup for s in self.tagged_sents[:self.fsize] \n",
    "                        for tup in s]\n",
    "        \n",
    "        self.vocabulary = (['_SWORD'] + # Special observation in STAG\n",
    "                      sorted(list(set([w.lower() for \n",
    "                                    s in self.sents for w in s]))) +\n",
    "                      ['_FWORD']) # Special observation in FTAG\n",
    "        self.M = len(self.vocabulary) # Number of distinct observations\n",
    "        \n",
    "        print(f'Vocabulary size {self.M}')\n",
    "    \n",
    "        self.tags = (['STAG'] + # Special start state\n",
    "                     sorted(list(set([p[1] for s in self.tagged_sents \n",
    "                                        for p in s]))) + \n",
    "                     ['FTAG']) # Special end state\n",
    "        self.N = len(self.tags) # Number of states\n",
    "    \n",
    "        self.w2i = dict(zip(self.vocabulary, \n",
    "                            range(len(self.vocabulary))))\n",
    "        self.t2i = dict(zip(self.tags, range(len(self.tags))))\n",
    "        \n",
    "        # Initialize a (transition probabilities)\n",
    "        #self.a = np.random.rand(self.N, self.N)\n",
    "        self.a = np.ones((self.N, self.N))\n",
    "        self.a[:,0] = 0.0\n",
    "        self.a = self.a/(self.a.sum(axis=1).reshape(\n",
    "            self.a.shape[0],1))\n",
    "        self.a[-1,:] = 0.0\n",
    "        \n",
    "        cfd = nltk.ConditionalFreqDist((word.lower(), tag) for \n",
    "                                       (word, tag)\n",
    "                                       in tagged_words)\n",
    "        \n",
    "        # Initialize b (emission probabilities)\n",
    "        #self.b = np.random.rand(self.N, self.M)*b_random_bound\n",
    "        self.b = np.ones((self.N, self.M))*b_bound\n",
    "        # STAG only produces _SWORD and FTAG only produces _FWORD\n",
    "        self.b[0,:] = 0.0\n",
    "        self.b[:,0] = 0.0\n",
    "        self.b[0,0] = 1.0\n",
    "        self.b[-1,:] = 0.0\n",
    "        self.b[:,-1] = 0.0\n",
    "        self.b[-1, -1] = 1.0\n",
    "        # Set non-trivial probability of a word being produced\n",
    "        # by its most common tag\n",
    "        for word in cfd.conditions():\n",
    "            tag = (cfd[word].most_common(1)[0][0])\n",
    "            self.setb(tag, word, 0.5)\n",
    "        self.b = self.b/(self.b.sum(axis=1).reshape(self.N,1))\n",
    "        \n",
    "        # Initialize an HMM on oset with a, b\n",
    "        self.hmm = HMMMulti(self.a.copy(), self.b.copy())\n",
    "        \n",
    "    def create_oset(self):\n",
    "        \n",
    "        self.oset = []\n",
    "        for ts in self.tagged_sents:\n",
    "            s, _ = split_tagged_sent(ts)\n",
    "            self.oset.append(self.s2o(s))\n",
    "        \n",
    "    def fit(self, max_iter=10):\n",
    "        \n",
    "        # Create oset\n",
    "        self.create_oset()\n",
    "            \n",
    "        # Fit the hmm \n",
    "        self.hmm.fit(self.oset, max_iter=max_iter)\n",
    "        \n",
    "    def test(self):\n",
    "        \n",
    "        for descrip, ts_set in [(\"fit set\", self.fit_tsents), \n",
    "                                (\"test set\", self.test_tsents)]:        \n",
    "            correct = 0\n",
    "            sum_accuracy = 0\n",
    "            for ts in ts_set:\n",
    "                s, tag_seq = split_tagged_sent(ts)\n",
    "                predicted, _ = self.decode(s)\n",
    "                a = accuracy(tag_seq, predicted)\n",
    "                if a == 1.0:\n",
    "                    correct += 1\n",
    "                sum_accuracy += a\n",
    "            avg_accuracy = sum_accuracy/len(ts_set)\n",
    "            correct_ratio = correct/len(ts_set)\n",
    "            print(f\"For {descrip}: Average accuracy = \"\n",
    "                  f\"{avg_accuracy:.3f},\" \n",
    "                  f\" Totally correct ratio = \" \n",
    "                  f\"{correct_ratio:.3f}\")\n",
    "              \n",
    "    def decode(self, sentence):\n",
    "        if sentence[0] != 0: # Weak check if converted\n",
    "            O = self.s2o(sentence)\n",
    "        else:\n",
    "            O = sentence.copy()\n",
    "        state_seq, prob = self.hmm.decode(O)\n",
    "        tag_seq = [self.tags[s] for s in state_seq]\n",
    "        return (tag_seq, prob)\n",
    "        \n",
    "    def setb(self, j, w, p):\n",
    "        self.b[self.t2i[j], self.w2i[w]] = p\n",
    "    def getb(self, j, w, b=None):\n",
    "        if b is None:\n",
    "            b = self.b\n",
    "        return b[self.t2i[j], self.w2i[w]]\n",
    "    def seta(self, i, j, p):\n",
    "        self.a[self.t2i[i], self.t2i[j]] = p\n",
    "    def geta(self, i, j, a=None):\n",
    "        if a is None:\n",
    "            a = self.a\n",
    "        return a[self.t2i[i], self.t2i[j]]\n",
    "    def s2o(self, sentence):\n",
    "        return([0] + \n",
    "               [self.w2i[w.lower()] for w in sentence] + \n",
    "               [self.w2i['_FWORD']])\n",
    "    def i2t(self, state_sequence):\n",
    "        return [self.tags[i] for i in state_sequence][1:-1]\n",
    "    \n",
    "    def print_b_unsorted(self, b=None):\n",
    "        if b is None:\n",
    "            b = self.b\n",
    "        for t in self.tags:\n",
    "            print(f'\\n\\n{t}', end='')\n",
    "            for w in self.vocabulary:\n",
    "                p = self.getb(t, w, b)\n",
    "                if p > 10**(-5):\n",
    "                    print(f' {w}[{p:.5f}]', end='')\n",
    "                    \n",
    "    def print_a_unsorted(self, a=None):\n",
    "        if a is None:\n",
    "            a = self.a\n",
    "        for t in self.tags:\n",
    "            print(f'\\n\\n{t}', end='')\n",
    "            for u in self.tags:\n",
    "                p = self.geta(t, u, a)\n",
    "                print(f' {u}[{p:.5f}]', end='')\n",
    "                \n",
    "    def print_b(self, b=None):\n",
    "        if b is None:\n",
    "            b = self.b\n",
    "        for t in self.tags:\n",
    "            ptups = zip(b[self.t2i[t],:],self.vocabulary)\n",
    "            ptups = sorted(ptups, key = lambda tup: tup[0],\n",
    "                          reverse = True) \n",
    "            print(f'\\n\\n{t}', end='')\n",
    "            for tup in ptups:\n",
    "                print(f' {tup[1]}[{tup[0]}]', end='')\n",
    "                \n",
    "    def print_a(self, a=None):\n",
    "        if a is None:\n",
    "            a = self.a\n",
    "        for t in self.tags:\n",
    "            ptups = zip(a[self.t2i[t],:],self.tags)\n",
    "            ptups = sorted(ptups, key = lambda tup: tup[0],\n",
    "                          reverse = True) \n",
    "            print(f'\\n\\n{t}', end='')\n",
    "            for tup in ptups:\n",
    "                print(f' {tup[1]}[{tup[0]}]', end='')\n",
    "                \n",
    "    def create_tagged_sent(self, sentence, tags):\n",
    "        if sentence[0] != '_SWORD':\n",
    "            s = (['_SWORD'] + \n",
    "                 [w.lower() for w in sentence] + \n",
    "                 ['_FWORD'])\n",
    "        else:\n",
    "            s = sentence\n",
    "        return list(zip(s, tags))\n",
    "    \n",
    "    def check_sent(self, sentence_number):\n",
    "        n = sentence_number\n",
    "        s = self.sents[n]\n",
    "        tag_seq = [tag for (word, tag) in \n",
    "                             self.tagged_sents[n]]\n",
    "        predicted_tag_seq = self.decode(s)[0][1:-1]\n",
    "        return list(zip(s, tag_seq, predicted_tag_seq))\n",
    "\n",
    "def split_tagged_sent(tagged_sent):\n",
    "    return list(zip(*tagged_sent))\n",
    "\n",
    "def accuracy(tag_seq, predicted_tag_seq):\n",
    "    if predicted_tag_seq[0] == 'STAG':\n",
    "        p = predicted_tag_seq[1:-1]\n",
    "    else:\n",
    "        p = predicted_tag_seq\n",
    "    return (sum(t == pt for (t, pt) in zip(tag_seq, p))/\n",
    "            len(tag_seq))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sents_example_1 = [\n",
    "    \"Janet will back the bill.\",\n",
    "    \"Richard will eat an apple.\",\n",
    "    \"Mary will run a mile.\",\n",
    "    \"Sam will write an essay.\",\n",
    "]\n",
    "\n",
    "start = time.time()\n",
    "pos = POS(sents_example_1)\n",
    "pos.fit(100)\n",
    "pos.test()\n",
    "pos.print_a(pos.hmm.a)\n",
    "pos.print_b(pos.hmm.b)\n",
    "finish = time.time()\n",
    "elapsed = finish - start\n",
    "print(\"time elapsed: \", elapsed)\n",
    "\n",
    "############################################################\n",
    "# PLEASE ANSWER THE FOLLOWING:  ############################\n",
    "############################################################\n",
    "# Did the HMM based POS tagger learn a correct POS\n",
    "# tagging model? If not, what might be the cause?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sents_example_2 = [\n",
    "    \"Janet will back the bill.\",\n",
    "    \"Janet likes music.\",\n",
    "    \"Richard will eat an apple.\",\n",
    "    \"Peter needs rest.\",\n",
    "    \"Mary will run a mile.\",\n",
    "    \"Joe hates TV.\",\n",
    "    \"Sam will write an essay.\",\n",
    "    \"Ramona loves cooking.\",\n",
    "]\n",
    "\n",
    "start = time.time()\n",
    "pos = POS(sents_example_2)\n",
    "pos.fit(100)\n",
    "pos.test()\n",
    "pos.print_a(pos.hmm.a)\n",
    "pos.print_b(pos.hmm.b)\n",
    "finish = time.time()\n",
    "elapsed = finish - start\n",
    "print(\"time elapsed: \", elapsed)\n",
    "\n",
    "############################################################\n",
    "# PLEASE ANSWER THE FOLLOWING:  ############################\n",
    "############################################################\n",
    "# Did the HMM based POS tagger learn a correct POS\n",
    "# tagging model? If not, what might be the cause?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from nltk import sent_tokenize\n",
    "\n",
    "with open(\"dev.txt\", \"r\") as fin:\n",
    "    text = fin.read()\n",
    "    \n",
    "sents_example_3 = sent_tokenize(text)\n",
    "\n",
    "start = time.time()\n",
    "pos = POS(sents_example_3)\n",
    "pos.fit(100)\n",
    "pos.test()\n",
    "pos.print_a(pos.hmm.a)\n",
    "pos.print_b(pos.hmm.b)\n",
    "finish = time.time()\n",
    "elapsed = finish - start\n",
    "print(\"time elapsed: \", elapsed)\n",
    "\n",
    "############################################################\n",
    "# PLEASE ANSWER THE FOLLOWING:  ############################\n",
    "############################################################\n",
    "# Did the HMM based POS tagger learn a correct POS\n",
    "# tagging model? If not, what might be the cause?"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
